
import torch
import torch.nn as nn

# -------------------- 模型定义 --------------------
class UNetDown(nn.Module):
    def __init__(self, in_size, out_size, normalize=True, dropout=0.0):
        super().__init__()
        layers = [nn.Conv2d(in_size, out_size, 4, 2, 1, bias=False)]
        if normalize:
            layers.append(nn.InstanceNorm2d(out_size))
        layers.append(nn.LeakyReLU(0.2))
        if dropout:
            layers.append(nn.Dropout2d(dropout))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)

class UNetUp(nn.Module):
    def __init__(self, in_size, out_size, dropout=0.0):
        super().__init__()
        self.model = nn.Sequential(
            nn.ConvTranspose2d(in_size, out_size, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(out_size),
            nn.ReLU(inplace=True),
            nn.Dropout2d(dropout) if dropout else nn.Identity()
        )

    def forward(self, x):
        return self.model(x)

# -------------------- CBAM模块定义 --------------------
class ChannelAttention(nn.Module):
    def __init__(self, channel, reduction=16):
        super().__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveMaxPool2d(1)
        self.mlp = nn.Sequential(
            nn.Linear(channel, channel // reduction),
            nn.ReLU(),
            nn.Linear(channel // reduction, channel)
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        B, C, _, _ = x.shape
        avg_out = self.mlp(self.avg_pool(x).view(B, C))
        max_out = self.mlp(self.max_pool(x).view(B, C))
        channel_weights = self.sigmoid(avg_out + max_out).view(B, C, 1, 1)
        return x * channel_weights

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size//2)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)
        max_out, _ = torch.max(x, dim=1, keepdim=True)
        concat = torch.cat([avg_out, max_out], dim=1)
        spatial_weights = self.sigmoid(self.conv(concat))
        return x * spatial_weights

class CBAM(nn.Module):
    def __init__(self, channel):
        super().__init__()
        self.channel_att = ChannelAttention(channel)
        self.spatial_att = SpatialAttention()

    def forward(self, x):
        x = self.channel_att(x)
        x = self.spatial_att(x)
        return x

# Generator with CBAM module
class GeneratorWithCBAM(nn.Module):
    def __init__(self):
        super().__init__()
        # 下采样层
        self.down1 = UNetDown(1, 64, normalize=False)        # 256 -> 128
        self.down2 = UNetDown(64, 128)                      # 128 -> 64
        self.down3 = UNetDown(128, 256)                     # 64 -> 32
        self.down4 = UNetDown(256, 512, dropout=0.3)       # 32 -> 16
        self.down5 = UNetDown(512, 512, dropout=0.3)        # 16 -> 8
        self.down6 = UNetDown(512, 512, dropout=0.3)        # 8 -> 4

        # 跳跃连接处的CBAM模块
        self.cbam1 = CBAM(64)    # 对应d1
        self.cbam2 = CBAM(128)   # 对应d2
        self.cbam3 = CBAM(256)   # 对应d3
        self.cbam4 = CBAM(512)   # 对应d4
        self.cbam5 = CBAM(512)   # 对应d5

        # 上采样层
        self.up1 = UNetUp(512, 512, dropout=0.3)            # 4 -> 8
        self.up2 = UNetUp(1024, 512, dropout=0.3)           # 8 -> 16
        self.up3 = UNetUp(1024, 256, dropout=0.3)           # 16 -> 32
        self.up4 = UNetUp(512, 128)                         # 32 -> 64
        self.up5 = UNetUp(256, 64)                          # 64 -> 128

        # 最终输出层
        self.up6 = nn.Sequential(
            nn.ConvTranspose2d(128, 64, 4, 2, 1),
            nn.InstanceNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 1, 3, padding=1),
            nn.Tanh()
        )

    def forward(self, x):
        # 下采样过程（编码器）
        d1 = self.down1(x)        # [B,64,128,128]
        d1_att = self.cbam1(d1)   # 对d1应用CBAM

        d2 = self.down2(d1)       # [B,128,64,64]
        d2_att = self.cbam2(d2)   # 对d2应用CBAM

        d3 = self.down3(d2)       # [B,256,32,32]
        d3_att = self.cbam3(d3)   # 对d3应用CBAM

        d4 = self.down4(d3)       # [B,512,16,16]
        d4_att = self.cbam4(d4)   # 对d4应用CBAM

        d5 = self.down5(d4)       # [B,512,8,8]
        d5_att = self.cbam5(d5)   # 对d5应用CBAM

        d6 = self.down6(d5)       # [B,512,4,4]

        # 上采样过程（解码器）
        u1 = self.up1(d6)
        u2 = self.up2(torch.cat([u1, d5_att], 1)) # 拼接d5_att
        u3 = self.up3(torch.cat([u2, d4_att], 1)) # 拼接d4_att
        u4 = self.up4(torch.cat([u3, d3_att], 1)) # 拼接d3_att
        u5 = self.up5(torch.cat([u4, d2_att], 1)) # 拼接d2_att
        u6 = self.up6(torch.cat([u5, d1_att], 1)) # 拼接d1_att

        return u6
    

class Discriminator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            # Layer 1: 256x256 -> 128x128
            nn.Conv2d(2, 64, 4, 2, 1),  # stride=2下采样
            nn.LeakyReLU(0.2, inplace=True),
            
            # Layer 2: 128x128 -> 64x64
            nn.Conv2d(64, 128, 4, 2, 1),
            nn.InstanceNorm2d(128),
            nn.LeakyReLU(0.2, inplace=True),
            
            # Layer 3: 64x64 -> 32x32
            nn.Conv2d(128, 256, 4, 2, 1),
            nn.InstanceNorm2d(256),
            nn.LeakyReLU(0.2, inplace=True),
            
            # Layer 4: 32x32 -> 16x16 (新增)
            nn.Conv2d(256, 512, 4, 2, 1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            
            # Layer 5: 16x16 -> 8x8 (新增)
            nn.Conv2d(512, 512, 4, 2, 1),
            nn.InstanceNorm2d(512),
            nn.LeakyReLU(0.2, inplace=True),
            
            # 最终输出层：8x8 -> 7x7（适合PatchGAN）
            nn.Conv2d(512, 1, 4, 1, 1)  # stride=1保持尺寸稳定
        )

    def forward(self, img_A, img_B):
        # 输入维度: (batch, 1, 256, 256) x 2
        # 拼接后: (batch, 2, 256, 256)
        concat = torch.cat([img_A, img_B], 1)
        return self.model(concat)  # 输出: (batch, 1, 7, 7)